#!/usr/bin/env python3
"""
将彩色图模型权重转换为灰度图模型权重

用法:
    python convert_color_to_gray.py \
        --input_checkpoint checkpoints/BSDS/10/10_model.pth \
        --output_checkpoint checkpoints_gray/BSDS/0/0_model.pth
"""

import argparse
import os
import torch
from collections import OrderedDict


def convert_color_to_gray_weights(color_checkpoint_path, gray_checkpoint_path, method='mean'):
    """
    将彩色图模型权重转换为灰度图模型权重
    
    Args:
        color_checkpoint_path: 彩色图模型checkpoint路径
        gray_checkpoint_path: 输出灰度图模型checkpoint路径
        method: 转换方法
            - 'mean': 取RGB三通道的平均值
            - 'luminance': 使用亮度加权 (0.299*R + 0.587*G + 0.114*B)
            - 'first': 仅使用第一个通道（B通道）
    """
    
    print(f"加载彩色图模型权重: {color_checkpoint_path}")
    assert os.path.isfile(color_checkpoint_path), f"文件不存在: {color_checkpoint_path}"
    
    checkpoint = torch.load(color_checkpoint_path, map_location='cpu')
    
    # 提取state_dict
    if isinstance(checkpoint, dict):
        if 'model' in checkpoint:
            state_dict = checkpoint['model']
        elif 'state_dict' in checkpoint:
            state_dict = checkpoint['state_dict']
        else:
            state_dict = checkpoint
    else:
        state_dict = checkpoint
    
    # 处理DataParallel前缀
    new_state_dict = OrderedDict()
    for key, value in state_dict.items():
        if key.startswith('module.'):
            key = key[7:]  # 移除 'module.' 前缀
        new_state_dict[key] = value
    state_dict = new_state_dict
    
    print(f"\n原始权重信息:")
    print(f"  总键数: {len(state_dict)}")
    
    # 查找第一层卷积权重
    first_conv_keys = [k for k in state_dict.keys() if 'block_1' in k and 'conv1.weight' in k]
    if not first_conv_keys:
        raise ValueError("未找到第一层卷积权重 (block_1.conv1.weight)")
    
    first_conv_key = first_conv_keys[0]
    conv1_weight = state_dict[first_conv_key]
    
    print(f"  第一层卷积键: {first_conv_key}")
    print(f"  第一层卷积形状: {conv1_weight.shape}")
    print(f"  期望形状: [32, 3, 3, 3] (out_channels, in_channels, H, W)")
    
    if conv1_weight.shape[1] != 3:
        raise ValueError(f"第一层输入通道数应为3，实际为{conv1_weight.shape[1]}")
    
    # 转换第一层卷积权重
    print(f"\n使用 '{method}' 方法转换第一层卷积...")
    
    if method == 'mean':
        # 简单平均三个通道
        gray_conv1_weight = conv1_weight.mean(dim=1, keepdim=True)
        print(f"  方法: 取三个通道的算术平均")
        
    elif method == 'luminance':
        # 使用人眼感知亮度权重
        # Y = 0.299*R + 0.587*G + 0.114*B (ITU-R BT.601标准)
        # OpenCV的BGR顺序: [B, G, R]
        weights = torch.tensor([0.114, 0.587, 0.299], dtype=conv1_weight.dtype)
        weights = weights.view(1, 3, 1, 1)  # reshape为可广播形状
        gray_conv1_weight = (conv1_weight * weights).sum(dim=1, keepdim=True)
        print(f"  方法: 亮度加权 (ITU-R BT.601)")
        print(f"  权重: B={weights[0,0,0,0]:.3f}, G={weights[0,1,0,0]:.3f}, R={weights[0,2,0,0]:.3f}")
        
    elif method == 'first':
        # 仅使用第一个通道（B通道）
        gray_conv1_weight = conv1_weight[:, 0:1, :, :]
        print(f"  方法: 仅使用B通道")
        
    else:
        raise ValueError(f"未知转换方法: {method}")
    
    print(f"  转换后形状: {gray_conv1_weight.shape}")
    print(f"  期望形状: [32, 1, 3, 3]")
    
    assert gray_conv1_weight.shape[1] == 1, "转换后应为单通道"
    
    # 更新state_dict
    state_dict[first_conv_key] = gray_conv1_weight
    
    # 构建新的checkpoint
    new_checkpoint = {
        'model': state_dict,
    }
    
    # 保留原checkpoint中的其他信息（如果存在）
    if isinstance(checkpoint, dict):
        if 'epoch' in checkpoint:
            new_checkpoint['epoch'] = 0  # 重置epoch，因为这是新模型
            print(f"\n注意: 原模型epoch={checkpoint['epoch']}，已重置为0")
        # 不复制optimizer状态，因为参数形状已改变
    
    # 保存
    os.makedirs(os.path.dirname(gray_checkpoint_path), exist_ok=True)
    torch.save(new_checkpoint, gray_checkpoint_path)
    
    print(f"\n✓ 灰度图模型权重已保存至: {gray_checkpoint_path}")
    
    # 验证转换结果
    print(f"\n验证转换结果...")
    from model import DexiNed
    
    try:
        # 加载灰度图模型
        model_gray = DexiNed(in_channels=1)
        model_gray.load_state_dict(state_dict, strict=True)
        
        # 测试前向传播
        test_input = torch.randn(1, 1, 320, 320)
        model_gray.eval()
        with torch.no_grad():
            outputs = model_gray(test_input)
        
        print(f"  ✓ 模型加载成功")
        print(f"  ✓ 前向传播成功，输出{len(outputs)}个特征图")
        print(f"  ✓ 输入形状: {test_input.shape}")
        print(f"  ✓ 输出形状: {outputs[-1].shape}")
        
    except Exception as e:
        print(f"  ✗ 验证失败: {e}")
        return False
    
    return True


def main():
    parser = argparse.ArgumentParser(description='将彩色图模型转换为灰度图模型')
    parser.add_argument('--input_checkpoint', type=str, required=True,
                        help='输入的彩色图模型checkpoint路径')
    parser.add_argument('--output_checkpoint', type=str, required=True,
                        help='输出的灰度图模型checkpoint路径')
    parser.add_argument('--method', type=str, default='mean',
                        choices=['mean', 'luminance', 'first'],
                        help='转换方法: mean(平均), luminance(亮度加权), first(仅B通道)')
    
    args = parser.parse_args()
    
    print("=" * 70)
    print("彩色图模型 -> 灰度图模型 权重转换工具")
    print("=" * 70)
    
    success = convert_color_to_gray_weights(
        args.input_checkpoint,
        args.output_checkpoint,
        args.method
    )
    
    if success:
        print("\n" + "=" * 70)
        print("✓ 转换成功！")
        print("\n现在可以使用转换后的模型进行灰度图训练或测试:")
        print(f"  python main.py --grayscale --resume --checkpoint_data {args.output_checkpoint}")
        print("=" * 70)
    else:
        print("\n" + "=" * 70)
        print("✗ 转换失败，请检查错误信息")
        print("=" * 70)
        return 1
    
    return 0


if __name__ == '__main__':
    import sys
    sys.exit(main())
